package leetcode.每日一题;

import org.junit.Test;

import java.util.*;

/**
 * @author :songyaozhang
 * @date :Created 2021/7/28 19:16
 * @description :https://leetcode-cn.com/problems/all-nodes-distance-k-in-binary-tree/
 */
public class 二叉树中所有距离为K的结点 {
    @Test
    public void test() {
        Solution solution = new Solution();

    }


    //Definition for a binary tree node.
    public class TreeNode {
        int val;
        TreeNode left;
        TreeNode right;
        TreeNode(int x) { val = x; }
    }

    class Solution {

        Map<Integer, TreeNode> parents = new HashMap<>();
        List<Integer> ans = new ArrayList<>();

        public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
            if (root == null || target == null) return null;
            findParents(root);
            findAns(target, null, 0, k);
            return ans;
        }

        private void findAns(TreeNode node, TreeNode from, int depth, int k) {
            if (node == null) return;
            if (depth == k) {
                ans.add(node.val);
                return;
            }
            if (node.left != from) {
                findAns(node.left, node, depth + 1, k);
            }
            if (node.right != from) {
                findAns(node.right, node, depth + 1, k);
            }
            if (parents.get(node.val) != from) {
                findAns(parents.get(node.val), node, depth + 1, k);
            }
        }

        private void findParents(TreeNode node) {
            if (node.left != null) {
                parents.put(node.left.val, node);
                findParents(node.left);
            }
            if (node.right != null) {
                parents.put(node.right.val, node);
                findParents(node.right);
            }
        }
    }


}





